package com.example.springbootlearn.netty.baseasr;

import cn.hutool.core.util.StrUtil;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import javax.annotation.PreDestroy;
import java.util.EnumMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;

@Service
public class AsrWsClientManager {

    private static final OkHttpClient CLIENT;
//    static ScheduledExecutorService executorService;

    static {
        Dispatcher dispatcher = new Dispatcher();
        dispatcher.setMaxRequests(80);
        CLIENT = new OkHttpClient.Builder().dispatcher(dispatcher).build();
    }

    private final Map<String, WebSocketInfo> sessionWebSocketMap;
    private final Map<CodeEnum, ScheduledExecutorService> executorServiceMap;
    private final Map<CodeEnum, AsrWsbBaseClient> clientMap;
    @Autowired
    private AsrConfig asrConfig;

    @Autowired
    public AsrWsClientManager(List<AsrWsbBaseClient> clients) {
        this.clientMap = new EnumMap<>(CodeEnum.class);
        this.executorServiceMap = new EnumMap<>(CodeEnum.class);
        this.sessionWebSocketMap = new ConcurrentHashMap<>();
        for (AsrWsbBaseClient client : clients) {
            AsrClient annotation = client.getClass().getAnnotation(AsrClient.class);
            if (annotation != null) {
                CodeEnum codeEnum = annotation.value();
                clientMap.put(codeEnum, client);
                if (!asrConfig.isDirectSend()) {
                    executorServiceMap.put(codeEnum, Executors.newScheduledThreadPool(asrConfig.getCorePoolSize())); // 可以根据需求调整线程池大小
                }
            }
        }
    }

    @PreDestroy
    public void shutdown() {
        for (ScheduledExecutorService executorService : executorServiceMap.values()) {
            executorService.shutdown();
            try {
                if (!executorService.awaitTermination(60, TimeUnit.SECONDS)) {
                    executorService.shutdownNow();
                }
            } catch (InterruptedException e) {
                executorService.shutdownNow();
            }
        }
    }

    private ScheduledFuture<?> startScheduleConsumer(CodeEnum codeEnum, String sessionId, ScheduledExecutorService executorService) {
        return executorService.scheduleAtFixedRate(() -> startConsumer(codeEnum, sessionId), asrConfig.getInitialDelay(), asrConfig.getPeriod(), TimeUnit.MILLISECONDS);
    }


    private void startConsumer(CodeEnum codeEnum, String sessionId) {
        WebSocketInfo webSocketSession = sessionWebSocketMap.get(sessionId);
        byte[] packet;
        if (webSocketSession != null) {
            packet = webSocketSession.poll();
            if (packet == null && asrConfig.getNeedHeartbeat()) {
                webSocketSession.sendHeartBeat();
            } else {
                webSocketSession.sendDataDirect(packet);
            }
        }
    }

    public WebSocketInfo createWebSocket(CodeEnum codeEnum, String sessionId) {
        AsrWsbBaseClient client = getClient(codeEnum);
        String wsUrl = client.getWsUrl(sessionId);
        Request request = new Request.Builder().url(wsUrl).build();
        WebSocket webSocket = CLIENT.newWebSocket(request, new WebSocketListener() {
            // 实现 WebSocketListener 的方法
        });
        String initialMessage = client.getInitialMessage(sessionId);
        if (StrUtil.isNotEmpty(initialMessage)) {
            webSocket.send(initialMessage);
        }

        WebSocketInfo webSocketInfo;
        ScheduledFuture<?> scheduledFuture = null;
        int needCount;
        if (asrConfig.isDirectSend()) {
            needCount = -1;
        } else {
            ScheduledExecutorService executorService = executorServiceMap.get(codeEnum);
            if (executorService != null) {
                scheduledFuture = startScheduleConsumer(codeEnum, sessionId, executorService);
            }
            needCount = asrConfig.getNeedCount();
        }
        webSocketInfo = client.createWebSocketInfo(webSocket, needCount, scheduledFuture);
        sessionWebSocketMap.put(sessionId, webSocketInfo);
        return webSocketInfo;
    }

    public void sendData(String sessionId, byte[] data) {
        WebSocketInfo webSocketInfo = sessionWebSocketMap.get(sessionId);
        if (webSocketInfo == null) {
            throw new IllegalStateException("No WebSocket found for session: " + sessionId);
        }
        webSocketInfo.sendDataByArray(data);
    }

    public void closeWebSocket(String sessionId) {
        WebSocketInfo webSocketInfo = sessionWebSocketMap.remove(sessionId);
        if (webSocketInfo != null) {
            webSocketInfo.close();
        }
    }

    private AsrWsbBaseClient getClient(CodeEnum codeEnum) {
        AsrWsbBaseClient client = clientMap.get(codeEnum);
        if (client == null) {
            throw new IllegalArgumentException("Unsupported client type: " + codeEnum);
        }
        return client;
    }
}